from ppgnss import gnss_utils
import xarray as xr
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.axes_grid1 import make_axes_locatable
plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']

width, height = gnss_utils.cm2inch(18), gnss_utils.cm2inch(8)

fontsize=10
filename = "/Users/lzhang/research/zihe/scripts/span5_comp_va_5min_snr1.txt"
"""
* MM    0.0015122  0.0004  117.01     25.05       5.2 * MM    0.0015122  0.0002   252.44     48.63       2.3   MM    0.0015122  0.0002    78.81     96.77      0.43 * MM    0.0015122     0.0734   301.40      5.71        65
"""
with open(filename, "r") as fread:
    lines = fread.readlines()
    freqs = []
    names = []
    amp_u_list, amp_v_list, amp_a_list = [], [], []
    snr_u_list, snr_v_list, snr_a_list = [], [], []
    amp_sl_list, snr_sl_list = [], []

    for line in lines[1:]:
        newline = line.replace("*", " ")
        fields = newline.split()
        name = fields[0]
        freq = float(fields[1])
        amp_u = float(fields[2])
        snr_u = float(fields[5])
        amp_v = float(fields[8])
        snr_v = float(fields[11])
        amp_a = float(fields[14])
        snr_a = float(fields[17])
        amp_sl = float(fields[20])
        snr_sl = float(fields[23])
        amp_u_list.append(amp_u)
        amp_v_list.append(amp_v)
        amp_a_list.append(amp_a)
        snr_u_list.append(snr_u)
        snr_v_list.append(snr_v)
        snr_a_list.append(snr_a)
        amp_sl_list.append(amp_sl)
        snr_sl_list.append(snr_sl)
        freqs.append(freq)
        names.append(name)

    data = [amp_u_list, snr_u_list, amp_v_list, snr_v_list, amp_a_list, snr_a_list]
    data = np.zeros((len(amp_u_list), 4, 2))
    data[:, 0, 0] = amp_u_list
    data[:, 0, 1] = snr_u_list
    data[:, 1, 0] = amp_v_list
    data[:, 1, 1] = snr_v_list
    data[:, 2, 0] = amp_a_list
    data[:, 2, 1] = snr_a_list
    data[:, 3, 0] = amp_sl_list
    data[:, 3, 1] = snr_sl_list

    xr_signal = xr.DataArray(data, dims=["name", "axis", "data"], coords=[names, ["u", "v", "a", "s"], ["amp", "snr"]])
    x = np.arange(len(xr_signal.coords["name"].values))
    y1 = [0, 1]
    yy1, xx1 = np.meshgrid(y1, x)

    snr_thred = 1
    # 创建图形和坐标轴
    fig, axes = plt.subplots(nrows=2, ncols=1, figsize=(width, height), gridspec_kw={'height_ratios': [1, 1.8]})
    z = xr_signal.loc[:, ["s"], "amp"].values
    mask = xr_signal.loc[:, ["s"], "snr"] < snr_thred
    z[mask] = np.nan
    print(xx1.shape, yy1.shape, z.shape)
    # axes[0].pcolormesh(np.arange(1, 3), np.arange(1, 11), matrix, cmap='viridis', edgecolors='k', linewidth=0.5)
    im = axes[0].pcolormesh(np.transpose(z), edgecolors="k", linewidth=0.2, cmap=plt.colormaps['magma_r'])
    axes[0].set_xticks(x+0.5)
    axes[0].set_xticklabels(names, fontsize=fontsize, rotation=270)
    axes[0].set_yticks([0.5])
    axes[0].set_yticklabels(["潮位"], fontsize=fontsize)
    axes[0].set_xlabel(u'潮汐分量', fontsize=fontsize)

    divider = make_axes_locatable(axes[0])
    cax = divider.new_vertical(size = '30%', pad = 0.3)
    fig.add_axes(cax)
    cbar = fig.colorbar(im, cax = cax, orientation = 'horizontal')
    ztklbls = np.arange(0, 2.6, 0.5)
    cbar.set_ticks(ztklbls)  # 根据刻度标签的数量设置刻度位置
    cbar.set_ticklabels(["%3.1f m" % _ for _ in ztklbls])

    ax = axes[1]
    y = [0, 1, 2]
    yy, xx = np.meshgrid(y, x)
    z = xr_signal.loc[:, ["u", "v", "a"], "amp"].values*1000
    mask = xr_signal.loc[:, ["u", "v", "a"], "snr"] < snr_thred
    z[mask] = np.nan

    plt.subplots_adjust(left=0.07, right=0.95, bottom=0.22, top=0.98, hspace=1)
    colors = cm.Greens(z / z.max())  # 使用viridis颜色映射，可以根据需要选择其他颜色映射
    # 绘制3D柱状图
    # 设置柱状图的尺寸
    dx,dy = 0.001, 0.1  # 在 x、y 轴方向上的宽度
    dz = z  # 在 z 轴方向上的高度
    im=ax.pcolormesh(xx, yy, z,  edgecolors="k", linewidth=0.2, cmap=plt.colormaps['RdYlGn_r'])
    ax.set_xticks(x)
    ax.set_xticklabels(names, fontsize=fontsize, rotation=270)
    ax.set_yticks(y)
    ax.set_yticklabels(["U", "Y", "X"], fontsize=fontsize)
    # 添加标签和标题
    ax.set_xlabel(u'潮汐分量', fontsize=fontsize)
    ax.set_ylabel(u'坐标分量', fontsize=fontsize)
    # ax.set_title('3D Bar Chart with Color Fill')
    # cbar = plt.colorbar(cax, orientation='horizontal')
    # add color bar below chart
    divider = make_axes_locatable(ax)
    cax = divider.new_vertical(size = '10%', pad = 0.3)
    fig.add_axes(cax)
    cbar = fig.colorbar(im, cax = cax, orientation = 'horizontal')
    ztklbls = np.arange(0, 1.5, 0.2)
    cbar.set_ticks(ztklbls)  # 根据刻度标签的数量设置刻度位置
    cbar.set_ticklabels(["%3.1f mm" % _ for _ in ztklbls])
    # 显示图形
    plt.show()